import torch
import matplotlib.pyplot as plt
import random


def create_data(w, b, data_num):  # 生成数据
    x = torch.normal(0, 1, (data_num, len(w)))
    y = torch.matmul(x, w) + b

    noise = torch.normal(0, 0.01, y.shape)  # 噪声要加到y上
    y += noise
    return x, y


num = 500
true_w = torch.tensor([8.1, 2, 2, 4])
true_b = 1.1
X, Y = create_data(true_w, true_b, num)

plt.scatter(X[:, 2], Y, 1)
plt.show()


def data_provider(data, label, batchsize):
    length = len(label)
    indices = list(range(length))
    # 不能按顺序取  把数据打乱
    # random.shuffle(indices)

    for each in range(0, length, batchsize):
        get_indices = indices[each: each + batchsize]
        get_data = data[get_indices]
        get_label = label[get_indices]

        yield get_data, get_label


batchsize = 16


for batch_x, batch_y in data_provider(X, Y, batchsize):
    print(batch_x, batch_y)
    break